use std::fmt;
use std::str::FromStr;

use aspect_codec::prelude::*;
use aspect_codec::{inst_semimanual_codec, inst_struct_codec, inst_wrapper_codec};

use openmls::prelude::{
    KeyPackage, KeyPackageIn, OpenMlsCrypto, ProtocolVersion, TlsDeserializeTrait,
};

use super::crypto::{self, IDSIG_PURPOSE_ATTESTATION};
use super::errors::*;
use super::sponsor;

/// Device identity ref.  Computed as the hash of the serialized pubkey.
#[derive(Copy, Clone, Hash, Eq, PartialEq, Ord, PartialOrd)]
pub struct DeviceIdent([u8; 32]);

impl DeviceIdent {
    /// Hashes a serialized ident pubkey to compute the device ident.
    pub fn from_ident_pubkey(ident_pk: &crypto::IdentPubkey) -> Self {
        let pkh = aspect_codec::util::compute_sha256(ident_pk).expect("signer: compute ident");
        Self(pkh)
    }

    /// Directly computes the device ident from the hash.
    pub fn from_hash(h: [u8; 32]) -> Self {
        Self(h)
    }

    /// Converts the inner has to a byte vec directly.
    pub fn to_vec(&self) -> Vec<u8> {
        self.0.to_vec()
    }
}

impl fmt::Debug for DeviceIdent {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let mut buf = [0u8; 64];
        hex::encode_to_slice(self.0, &mut buf).expect("device: encode device ident");
        let ss = unsafe { std::str::from_utf8_unchecked(&buf) };
        f.write_str(ss)
    }
}

impl fmt::Display for DeviceIdent {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        <Self as fmt::Debug>::fmt(self, f)
    }
}

impl AsRef<[u8; 32]> for DeviceIdent {
    fn as_ref(&self) -> &[u8; 32] {
        &self.0
    }
}

impl TryFrom<&[u8]> for DeviceIdent {
    type Error = Error;

    fn try_from(value: &[u8]) -> Result<Self, Self::Error> {
        if value.len() != 32 {
            return Err(Error::MalformedIdent);
        }

        let mut buf = [0; 32];
        buf.copy_from_slice(&value[..32]);
        Ok(Self::from_hash(buf))
    }
}

impl FromStr for DeviceIdent {
    type Err = Error;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        if s.len() != 64 {
            return Err(Error::MalformedIdent);
        }

        let mut buf = [0; 32];
        hex::decode_to_slice(s, &mut buf).map_err(|_| Error::MalformedIdent)?;
        Ok(Self::try_from(&buf[..])?)
    }
}

inst_wrapper_codec!(DeviceIdent => [u8; 32]);

/// Initial announcement of a device.
#[derive(Clone, Debug)]
pub struct AnnounceData {
    birth_time: u64,
    ident_pk: crypto::IdentPubkey,
    short_name: String,
    sponsorship: Option<sponsor::Attestation>,
}

impl AnnounceData {
    pub fn new(
        birth_time: u64,
        ident_pk: crypto::IdentPubkey,
        short_name: String,
        sponsorship: Option<sponsor::Attestation>,
    ) -> Self {
        Self {
            birth_time,
            ident_pk,
            short_name,
            sponsorship,
        }
    }

    pub fn ident_pk(&self) -> &crypto::IdentPubkey {
        &self.ident_pk
    }

    pub fn get_ident(&self) -> DeviceIdent {
        DeviceIdent::from_ident_pubkey(&self.ident_pk)
    }
}

inst_struct_codec! {
    AnnounceData {
        birth_time: u64,
        ident_pk: crypto::IdentPubkey,
        short_name: String,
        sponsorship: Option<sponsor::Attestation>,
    }
}

/// Signed announcement package.
#[derive(Clone, Debug)]
pub struct AnnPackage {
    /// Signature (computed from `id,H(data)`).
    sig: crypto::Signature,

    /// Serialized announcement data.
    ann_data: Vec<u8>,
}

impl AnnPackage {
    pub fn new(sig: crypto::Signature, ann_data: Vec<u8>) -> Self {
        Self { sig, ann_data }
    }

    pub fn sig(&self) -> &crypto::Signature {
        &self.sig
    }

    pub fn announcement_raw(&self) -> &[u8] {
        &self.ann_data
    }

    /// Tries to decode the announcement data as an `AnnounceData` structure.
    pub fn try_decode(&self) -> Result<AnnounceData, CodecError> {
        aspect_codec::util::decode(self.announcement_raw())
    }
}

inst_struct_codec! {
    AnnPackage {
        sig: crypto::Signature,
        ann_data: Vec<u8>,
    }
}

/// Describes shortlived state about a device.
#[derive(Clone, Debug)]
pub struct UpdateData {
    /// Update timestamp.
    timestamp: u64,

    /// Milliseconds to trust and respect this update for, unless one with a higher seqno is seen.
    lifetime: u64,

    /// Key used to encrypt signalling messages to the device.
    inbox_key: crypto::InboxPubkey,

    /// Extra arbitrary data about the device identity, maybe used to explain an update?
    comment: String,

    /// MLS key packages used to initiate sessions.
    prekeys: Vec<KeypackageContainer>,
}

impl UpdateData {
    pub fn new(
        timestamp: u64,
        lifetime: u64,
        inbox_key: crypto::InboxPubkey,
        comment: String,
        prekeys: Vec<KeypackageContainer>,
    ) -> Self {
        Self {
            timestamp,
            lifetime,
            inbox_key,
            comment,
            prekeys,
        }
    }
}

inst_semimanual_codec! {
    UpdateData {
        timestamp field u64,
        lifetime field u64,
        inbox_key field crypto::InboxPubkey,
        comment field String,
        prekeys vec KeypackageContainer
    }
}

/// Signed device identity record.
#[derive(Clone, Debug)]
pub struct UpdatePackage {
    /// Identity.
    id: DeviceIdent,

    /// Record seqno (incrementing from 0).
    seqno: u32,

    /// Signature (message computed from `id,seqno,H(data)`).
    sig: crypto::Signature,

    /// Serialized record.
    update_data: Vec<u8>,
}

impl UpdatePackage {
    pub fn new(id: DeviceIdent, seqno: u32, sig: crypto::Signature, update_data: Vec<u8>) -> Self {
        Self {
            id,
            seqno,
            sig,
            update_data,
        }
    }

    pub fn ident(&self) -> &DeviceIdent {
        &self.id
    }

    pub fn seqno(&self) -> u32 {
        self.seqno
    }

    pub fn sig(&self) -> &crypto::Signature {
        &self.sig
    }

    pub fn record_raw(&self) -> &[u8] {
        &self.update_data
    }
}

inst_struct_codec! {
    UpdatePackage {
        id: DeviceIdent,
        seqno: u32,
        sig: crypto::Signature,
        update_data: Vec<u8>,
    }
}

/// Combined DIR data, both its announcement and a recent update.
#[derive(Clone, Debug)]
pub struct DirPackage {
    ann: AnnPackage,
    last_update: UpdatePackage,
}

impl DirPackage {
    pub fn new(ann: AnnPackage, last_update: UpdatePackage) -> Self {
        Self { ann, last_update }
    }

    pub fn announcement(&self) -> &AnnPackage {
        &self.ann
    }

    pub fn last_update(&self) -> &UpdatePackage {
        &self.last_update
    }

    /// Returns the sequence number on the update package.
    pub fn seqno(&self) -> u32 {
        self.last_update().seqno()
    }
}

inst_struct_codec! {
    DirPackage {
        ann: AnnPackage,
        last_update: UpdatePackage,
    }
}

/// Report used to report a device making two conflicting announces/updates that
/// would break lots of things.  When a client sees this it should try to avoid
/// interacting with the device in question and use this as a gravestone it can
/// share with others.
#[derive(Clone, Debug)]
pub struct EquivocationReport {
    /// The device that equivocated.
    id: DeviceIdent,

    /// The seqno that was being equivocated, or `None` if it was an announcement.
    seqno: Option<u32>,

    /// The first (1) record the creator of this report saw.
    record_1: Vec<u8>,

    /// Its signature (1).
    sig_1: crypto::Signature,

    /// The second (2) record the creator of this report saw.
    record_2: Vec<u8>,

    /// Its signature (2).
    sig_2: crypto::Signature,
}

impl EquivocationReport {
    /// Compares two DIR packages and constructs an equivocation report.  Does
    /// not check signatures, assumes they are independently well-formed.
    pub fn from_pkgs(a: &DirPackage, b: &DirPackage) -> Result<Self, Error> {
        if a.last_update().ident() != b.last_update().ident() {
            return Err(Error::EquivUnrelated(
                *a.last_update().ident(),
                *b.last_update().ident(),
            ));
        }

        let ident = a.last_update().ident();

        // Try constructing from announcements since that's really bad.
        match Self::from_announcements(*ident, a.announcement(), b.announcement()) {
            Ok(eqv) => return Ok(eqv),

            // If this didn't work, move on.
            Err(_) => {}
        }

        // Now try constructing from the last update packages.
        Self::from_updates(a.last_update(), b.last_update())
    }

    /// Constructs an equivocation report from two different DIR update
    /// packages.  If they're the same, returns error.  Does not validate that
    /// the signatures are correct.
    pub fn from_announcements(
        id: DeviceIdent,
        a: &AnnPackage,
        b: &AnnPackage,
    ) -> Result<Self, Error> {
        // If the two payloads are the same they're the same DIR.
        if a.ann_data == a.ann_data {
            return Err(Error::EquivSameDir);
        }

        Ok(Self {
            id,
            seqno: None,
            record_1: a.ann_data.clone(),
            sig_1: a.sig.clone(),
            record_2: b.ann_data.clone(),
            sig_2: b.sig.clone(),
        })
    }

    /// Constructs an equivocation report from two different DIR update
    /// packages.  If they're the same, returns error.  Does not validate that
    /// the signatures are correct.
    pub fn from_updates(a: &UpdatePackage, b: &UpdatePackage) -> Result<Self, Error> {
        // If the seqnos are different then they're obviously unrelated.  Catch
        // this explicitly.
        if a.seqno != b.seqno {
            return Err(Error::EquivDifferentSeqnos);
        }

        // If the two payloads are the same they're the same DIR.
        if a.update_data == b.update_data {
            return Err(Error::EquivSameDir);
        }

        Ok(Self {
            id: a.id,
            seqno: Some(a.seqno),
            record_1: a.update_data.clone(),
            sig_1: a.sig.clone(),
            record_2: b.update_data.clone(),
            sig_2: b.sig.clone(),
        })
    }

    /// If the equivocation is from an update, returns its sequence number.
    pub fn seqno(&self) -> Option<u32> {
        self.seqno
    }
}

inst_struct_codec! {
    EquivocationReport {
        id: DeviceIdent,
        seqno: Option<u32>,
        record_1: Vec<u8>,
        sig_1: crypto::Signature,
        record_2: Vec<u8>,
        sig_2: crypto::Signature
    }
}

impl EquivocationReport {
    pub fn ident(&self) -> &DeviceIdent {
        &self.id
    }
}

/// Wraps a serialized MLS key package.  Not guaranteed to be valid.
#[derive(Clone, Debug)]
pub struct KeypackageContainer(Vec<u8>);

impl KeypackageContainer {
    /// Parses and validates the key package.
    pub fn parse_as_proto(&self, crypto: &impl OpenMlsCrypto) -> Result<KeyPackage, Error> {
        let ver = ProtocolVersion::Mls10;
        let pki = KeyPackageIn::tls_deserialize_exact(&self.0)?;
        Ok(pki.validate(crypto, ver)?)
    }
}

inst_wrapper_codec!(KeypackageContainer => Vec<u8>);

/// Used to uniquely refer to a prekey assuming packages were published
/// appropriately.
#[derive(Copy, Clone, Debug)]
pub struct DevPrekeyRef {
    /// The DID.
    ident: DeviceIdent,

    /// Device identity record seqno.
    seqno: u32,

    /// Which key package seqno.
    index: u16,
}

inst_struct_codec! {
    DevPrekeyRef {
        ident: DeviceIdent,
        seqno: u32,
        index: u16,
    }
}

/// Attestation payload that can be sent over the network to a challenger that
/// provides enough information to prove ownership of the device signature.
#[derive(Clone, Debug)]
pub struct DeviceAttestation {
    ann: AnnPackage,
    challenge_sig: crypto::Signature,
}

impl DeviceAttestation {
    pub fn new(ann: AnnPackage, challenge_sig: crypto::Signature) -> Self {
        Self { ann, challenge_sig }
    }

    pub fn ann_package(&self) -> &AnnPackage {
        &self.ann
    }

    pub fn challenge_sig(&self) -> &crypto::Signature {
        &self.challenge_sig
    }

    pub fn is_ann_well_formed(&self) -> bool {
        self.ann.try_decode().is_ok()
    }

    pub fn get_ident(&self) -> Option<DeviceIdent> {
        self.ann.try_decode().ok().map(|ann| ann.get_ident())
    }
}

pub fn verify_attestation(attest: &DeviceAttestation, chal: &[u8]) -> bool {
    let Ok(ann) = attest.ann_package().try_decode() else {
        return false;
    };

    let pk = ann.ident_pk();
    if !crypto::verify_ident_msg(attest.challenge_sig(), pk, IDSIG_PURPOSE_ATTESTATION, chal) {
        return false;
    }

    true
}

inst_struct_codec! {
    DeviceAttestation {
        ann: AnnPackage,
        challenge_sig: crypto::Signature,
    }
}
